# coding=utf-8

import re
import hashlib
import asyncio
from collections.abc import Generator
from typing import Any, Dict

from dify_plugin import Tool
from dify_plugin.entities.tool import ToolInvokeMessage

from .call_vl_model import call_vl_model

class DocImageAltTool(Tool):

    ERROR_CONTENT = "无法获取图片描述"

    def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
        """
        用于将图片替换为描述图片内容的文字。

        Args:
            tool_parameters: 一个包含工具输入参数的字典:
                - p_content (str): 需要将图片替换为文字描述的Markdown文本内容

        Yields:
            ToolInvokeMessage: 包含完整Markdown文本信息，其中图片链接中增加了对图片的文字描述。

        Raises:
            Exception: 如果文本生成失败，抛出异常。
        """
        try:
            self.base_url = self.runtime.credentials["base_url"]
        except KeyError:
            raise Exception("Base URL 未配置或无效。请在插件设置中提供。")
        try:
            self.access_token = self.runtime.credentials["access_token"]
        except KeyError:
            raise Exception("Access Token 未配置或无效。请在插件设置中提供。")

        content = tool_parameters.get("p_content", "")
        if not content:
            raise Exception("MarkDown文本内容不能为空。")

        # 获取最大并发数参数，默认为5
        max_concurrency = tool_parameters.get("max_concurrency", 5)
        if not isinstance(max_concurrency, int) or max_concurrency <= 0:
            max_concurrency = 5
        print(f"[DocImageAltTool] 开始处理Markdown文本内容，最大并发数: {max_concurrency}")
        print(f"[DocImageAltTool] Base URL: {self.base_url}, Access Token: {10 * '*' + self.access_token[10:-1]}")
        # 处理Markdown文本，提取图片URL并生成唯一标识符替换alt文本
        result = self._process_markdown_images(content)
        processed_content = result['processed_content']
        url_to_id = result['url_to_id']

        # 获取图片描述并替换到Markdown文本中
        processed_content, processed_result = asyncio.run(self._replace_markdown_images(url_to_id=url_to_id, content=processed_content, max_concurrency=max_concurrency))

        # yield self.create_json_message({
        #     "result": processed_content
        # })
        print(f"[DocImageAltTool] 已完成Markdown文本处理")
        yield self.create_variable_message("content", processed_content)
        yield self.create_variable_message("result", processed_result)

    def _process_markdown_images(self, content: str) -> Dict[str, Any]:
        """
        处理Markdown文本中的图片，为每个URL生成唯一标识符并替换alt文本

        Args:
            content: 原始Markdown文本

        Returns:
            处理后的Markdown文本和图片URL字典
        """
        # 用于存储URL到标识符的映射，确保相同URL使用相同标识符
        url_to_id = {}

        # 匹配Markdown图片的正则表达式：![alt](url)
        pattern = r'!\[([^\]]*)\]\(([^)]+)\)'

        def replace_func(match):
            alt_text = match.group(1)
            url = match.group(2)

            # 为URL生成唯一标识符
            if url in url_to_id:
                identifier = url_to_id[url]
            else:
                # 使用URL的hash值生成唯一标识符
                url_hash = hashlib.md5(url.encode('utf-8')).hexdigest()
                identifier = f"img_{url_hash}"
                url_to_id[url] = identifier

            # 返回替换后的图片标记
            return f"![{identifier}]({url})"

        # 替换所有匹配的图片标记
        processed_content = re.sub(pattern, replace_func, content)
        return {
            "processed_content": processed_content,
            "url_to_id": url_to_id,
        }

    async def _replace_markdown_images(self, url_to_id: Dict, content: str, max_concurrency: int = 5) -> tuple[str, Dict[str, int]]:
        """
        使用协程并发处理图片描述替换

        Args:
            url_to_id: URL到标识符的映射字典
            content: 处理后的Markdown内容
            max_concurrency: 最大并发数，默认为5

        Returns:
            tuple: 包含两个元素的元组:
                - str: 替换图片描述后的Markdown内容
                - Dict[str, int]: 处理结果统计，包含'all'(总数)、'success'(成功数)、'failure'(失败数)
        """
        # 创建信号量控制并发数
        semaphore = asyncio.Semaphore(max_concurrency)

        # 创建任务列表
        tasks = []
        identifiers = []
        urls = []

        # 为每个URL创建异步任务
        for url, identifier in url_to_id.items():
            task = self._process_single_image(url, identifier, semaphore)
            tasks.append(task)
            identifiers.append(identifier)
            urls.append(url)

        # 并发执行所有任务
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # 处理结果并替换内容
        for i, identifier in enumerate(identifiers):
            result = results[i]

            # 处理异常情况
            if isinstance(result, Exception):
                # 获取对应的URL用于错误信息
                url = urls[i]
                print(f"[DocImageAltTool] 处理图片 {url} 时发生错误: {str(result)}")
                result = self.ERROR_CONTENT

            # 替换内容中的标识符
            content = content.replace(identifier, result)

        # 统计图片处理结果
        processed_result = {
            'all': len(urls),
            'failure': content.count(self.ERROR_CONTENT),
            'success': len(urls) - content.count(self.ERROR_CONTENT),
        }

        return content, processed_result

    async def _process_single_image(self, url: str, identifier: str, semaphore: asyncio.Semaphore) -> str:
        """
        处理单个图片的异步方法

        Args:
            url: 图片URL
            identifier: 图片标识符
            semaphore: 用于控制并发的信号量

        Returns:
            图片描述文本
        """
        async with semaphore:  # 控制并发数
            print(f"[DocImageAltTool] 开始处理图片：{url}")
            try:
                # 在线程池中运行阻塞的call_vl_model函数
                loop = asyncio.get_running_loop()
                response = await loop.run_in_executor(None, call_vl_model, url, self.base_url, self.access_token)
                if response:
                    response = response.replace("\n", " ").replace("\r", " ").replace("\t", " ")
                else:
                    response = self.ERROR_CONTENT
                return response
            except Exception as e:
                print(f"[DocImageAltTool] 处理图片 {url} 时发生错误: {str(e)}")
                return self.ERROR_CONTENT